/***********************************************************************************************************
* TNA - The Network Auditor
* Copyright © 2012  Jeff Scaparra -- SPAWAR Systems Center LANT
*
* This file is a part of TNA.
************************************************************************************************************/

#include "tcpserviceportdiscoverer.h"
#include "helper.h"
#include <QDateTime>

TCPServicePortDiscoverer::TCPServicePortDiscoverer(): filename_("serviceLog.csv"), file_(NULL)
{
  logToFile_ = false;
  logToStdOut_ = true;
  logToDB_ = false;
  useUTC_ = true;
  timestampFormat_ = "M-d-yy hh:mm:ss";
  IPv4Address v4network( "0.0.0.0" );
  IPv4Address v4netmask( "0.0.0.0" );
  QPair< IPv4Address, IPv4Address > v4netblock( v4network, v4netmask );
  v4Netblocks_.push_back( v4netblock );
  IPv6Address v6network( "::" );
  IPv6Address v6netmask( "::" );
  QPair< IPv6Address, IPv6Address > v6netblock( v6network, v6netmask );
  v6Netblocks_.push_back( v6netblock );
}

TCPServicePortDiscoverer::~TCPServicePortDiscoverer()
{
  if( file_ )
    delete file_;
}

void TCPServicePortDiscoverer::processPacket(const Packet &p)
{
  bool destinationInNetblock = false;
  bool sourceInNetblock = false;

  std::string sourceString;
  std::string destinationString;

  if( p.transIs<TCP>() ) //is it tcp
  {
    TCP tcp = p.getTrans<TCP>();
    uint16_t sport = tcp.sourcePort();
    uint16_t dport = tcp.destinationPort();
    uint16_t seqnum = tcp.sequenceNumber();

    if( p.inetIs<IPv4>() ) //is IPv4
    {
      IPv4 ip = p.getInet<IPv4>();
      IPv4Address source = ip.sourceAddress();
      IPv4Address destination = ip.destinationAddress();
      destinationInNetblock = isInV4Netblocks(destination);
      sourceInNetblock = isInV4Netblocks(source);
      sourceString = source.toString() +"#" +ss_itoa<uint16_t>( sport );
      destinationString = destination.toString() +"#" + ss_itoa<uint16_t>(dport);
    }
    if( p.inetIs<IPv6>() )
    {
      IPv6 ip = p.getInet<IPv6>();
      IPv6Address source = ip.sourceAddress();
      IPv6Address destination = ip.destinationAddress();
      destinationInNetblock = isInV6Netblocks(destination);
      sourceInNetblock = isInV6Netblocks(source);
      sourceString = source.toString() +"#" +ss_itoa<uint16_t>( sport );
      destinationString = destination.toString() +"#" + ss_itoa<uint16_t>(dport);
    }

    if( tcp.SYN_Flag() == true )
    {
      QDateTime time;
      if( useUTC_ )
      {
        time = QDateTime::currentDateTimeUtc();
      }
      else
      {
        time = QDateTime::currentDateTime();
      }
      ServiceQueueData data;
      data.timestamp = time;
      data.sequenceNum = seqnum;

      if( tcp.ACK_Flag() == false && tcp.FIN_Flag() == false && tcp.RST_Flag() == false && destinationInNetblock )
      { //new possible connection
        synQueue_[destinationString][sourceString] = data;
      }
      else if( tcp.ACK_Flag() == true && tcp.FIN_Flag() == false && tcp.RST_Flag() == false && sourceInNetblock )
      { //if in synQueue then connection
        if( synQueue_.contains( sourceString ) )
          if( synQueue_[sourceString].contains( destinationString ) )
          {
            //We have a connection
            log( destinationString + ", " + sourceString + ", " + time.toString( "M-d-yy hh:mm:ss" ).toStdString() + "\n");
            removeConnection( sourceString, destinationString );
          }
      }
    }
    else if( tcp.RST_Flag() == true )
    {
      removeConnection( sourceString, destinationString );
      removeConnection( destinationString, sourceString );
    }
  }
}


void TCPServicePortDiscoverer::removeConnection(const std::string &sourceString, const std::string &destinationString)
{
  if( synQueue_.contains( sourceString ) )
  {
    synQueue_[sourceString].remove(destinationString);
    if( synQueue_[sourceString].size() == 0 )
      synQueue_.remove(sourceString);
  }
}

void TCPServicePortDiscoverer::setUTC(const bool &enable)
{
  useUTC_ = enable;
}

bool TCPServicePortDiscoverer::utc() const
{
  return useUTC_;
}

void TCPServicePortDiscoverer::log(const std::string &message)
{
  if( logToStdOut_ )
  {
    std::cout << message << "\n";
  }
  if( logToFile_ )
  {
    if( file_ == NULL )
    { //create file pointer
      file_ = new std::ofstream();
      file_->open(filename_.c_str());
    }
    *file_ << message << "\n";
   }
  if( logToDB_ )
  {
    //TODO: Add DB LOGGING
  }
}

void TCPServicePortDiscoverer::setLogToFile(const bool &enable)
{
  logToFile_ = enable;
}

void TCPServicePortDiscoverer::setLogToStdOut(const bool &enable)
{
  logToStdOut_ = enable;
}

bool TCPServicePortDiscoverer::logToFile() const
{
  return logToFile_;
}

bool TCPServicePortDiscoverer::logToStdOut() const
{
  return logToStdOut_;
}

void TCPServicePortDiscoverer::setLogFile(const std::string &filename)
{
  if( file_ )
  {
    file_->close();
    delete file_;
    file_ = NULL;
  }
  filename_ = filename;
}

std::string TCPServicePortDiscoverer::logFile() const
{
  return filename_;
}

void TCPServicePortDiscoverer::setTimestampString(const std::string &format)
{
  timestampFormat_ = format;
}

std::string TCPServicePortDiscoverer::timestampFormat() const
{
  return timestampFormat_;
}

QVector< QPair< IPv4Address, IPv4Address > > TCPServicePortDiscoverer::v4Netblocks() const
{
  return v4Netblocks_;
}

void TCPServicePortDiscoverer::clearV4Netblocks()
{
  v4Netblocks_.clear();
}

void TCPServicePortDiscoverer::addV4Netblock(const IPv4Address &network, const IPv4Address &netmask)
{
  QPair< IPv4Address, IPv4Address > netblock( network, netmask );
  v4Netblocks_.push_back( netblock );
}

QVector< QPair< IPv6Address, IPv6Address > > TCPServicePortDiscoverer::v6Netblocks() const
{
  return v6Netblocks_;
}

void TCPServicePortDiscoverer::clearV6Netblocks()
{
  v6Netblocks_.clear();
}

void TCPServicePortDiscoverer::addV6Netblock(const IPv6Address &network, const IPv6Address &netmask)
{
  QPair< IPv6Address, IPv6Address > netblock( network, netmask );
  v6Netblocks_.push_back( netblock );
}

bool TCPServicePortDiscoverer::isInV4Netblocks(const IPv4Address &address) const
{
  QVectorIterator< QPair< IPv4Address, IPv4Address> > itr( v4Netblocks_ );
  while( itr.hasNext() )
  {
    QPair< IPv4Address, IPv4Address > network = itr.next();
    bool found = (network.first & network.second) == (address & network.second);
    if( found )
      return true;
  }
  return false;
}

bool TCPServicePortDiscoverer::isInV6Netblocks(const IPv6Address &address) const
{
  QVectorIterator< QPair< IPv6Address, IPv6Address> > itr( v6Netblocks_ );
  while( itr.hasNext() )
  {
    QPair< IPv6Address, IPv6Address > network = itr.next();
    bool found = (network.first & network.second) == (address & network.second);
    if( found )
      return true;
  }
  return false;
}
